# 模型预测

import os

import paddle
from paddleseg.core import predict
from paddleseg.models import DeepLabV3, ResNet50_vd

from src.color.color import trans2color2list
from src.config.config import DATA_ROOT, TEST_PATH, NUM_CLASSES, TEST_SAVE_DIR, MODEL_PATH
from src.util.dataset import test_transforms_compose

# 预测测试集
model = DeepLabV3(
    num_classes=NUM_CLASSES,
    backbone=ResNet50_vd(),  # currently support Resnet50_vd/Resnet101_vd/Xception65.
    pretrained=None
)

# 设置模型参数
if MODEL_PATH:
    para_state_dict = paddle.load(MODEL_PATH)
    model.set_dict(para_state_dict)
    print('Loaded trained params of model successfully')
else:
    raise ValueError('The model_path is wrong: {}'.format(MODEL_PATH))


def get_image_list(data_root, test_path):
    """
    Get image list from the test path
    """

    image_list = []
    with open(test_path, 'r') as file:
        lines = file.readlines()
        for line in lines:
            image_path = os.path.join(data_root, line.split(' ')[0])
            if os.path.isfile(image_path):
                image_list.append(image_path)

    if len(image_list) == 0:
        raise RuntimeError('There are not image file in `--image_path`')

    return image_list


# 测试集图像文件目录、文件名列表
image_dir = os.path.join(DATA_ROOT, 'images')
image_list = get_image_list(DATA_ROOT, TEST_PATH)
# print(image_list)
# print(image_dir)


predict(
    model,
    model_path=MODEL_PATH,  # 模型路径
    transforms=test_transforms_compose,  # transform.Compose, 对输入图像进行预处理
    image_list=image_list,  # list, 待预测的图像路径列表。
    image_dir=image_dir,  # str, 待预测的图片所在目录
    save_dir=TEST_SAVE_DIR,  # str, 结果输出路径
    custom_color=trans2color2list()  # 设置自定义分割预测颜色，len(custom_color) = 3 * 像素种类
)

# ```
# 预测结果保存位置
# ./data/images
# ./output/results/added_prediction
# ./output/results/pseudo_color_prediction
# ```
# | 原始图片 | 叠加预测结果 | 预测结果 |
# | -------- | -------- | -------- |
